package com.lzq.mq.broker.rebalance;

import com.lzq.mq.broker.cache.CommonCache;
import com.lzq.mq.broker.rebalance.strategy.IReBalanceStrategy;
import com.lzq.mq.broker.rebalance.strategy.ReBalanceInfo;
import com.lzq.mq.broker.rebalance.strategy.impl.RandomReBalanceStrategyImpl;
import com.lzq.mq.broker.rebalance.strategy.impl.RangeReBalanceStrategyImpl;
import com.lzq.mq.common.utils.CollectionUtils;
import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * 消费者实例池
 * 存储消费者实例信息
 */
@Slf4j
public class RebalanceConsumerInstancePool {

    /**
     * 存储订阅主题下的消费者实例应映射
     * key: topic
     * value: List<ConsumerInstance>
     */
    private final Map<String, List<ConsumerInstance>> consumeInstanceMap = new ConcurrentHashMap<>();

    /**
     * 重平衡策略, 存储重平衡策略名称到策略执行对象的映射
     * key: strategyName
     * value: IReBalanceStrategy
     */
    private static final Map<String, IReBalanceStrategy> reBalanceStrategyMap = new HashMap<>();

    /**
     *
     */
    private final ReBalanceInfo reBalanceInfo = new ReBalanceInfo();

    static {
        reBalanceStrategyMap.put("random", new RandomReBalanceStrategyImpl());
        reBalanceStrategyMap.put("range", new RangeReBalanceStrategyImpl());
    }

    /**
     * 加入到池中
     * 如果存在则不做处理
     */
    public void addInstancePool(ConsumerInstance consumerInstance) {
        synchronized (this) {
            String topic = consumerInstance.getTopic();
            List<ConsumerInstance> consumerInstanceList =
                    consumeInstanceMap.computeIfAbsent(topic, k -> new ArrayList<>());

            for (ConsumerInstance instance : consumerInstanceList) {
                if (instance.getConsumerReqId().equals(consumerInstance.getConsumerReqId())) {
                    return;     // 存在则不做处理
                }
            }

            consumerInstanceList.add(consumerInstance);

            Set<String> consumerGroupSet =
                    reBalanceInfo.getChangeConsumerGroupMap().computeIfAbsent(topic, k -> new HashSet<>());
            consumerGroupSet.add(consumerInstance.getConsumeGroup());
        }
    }

    /**
     * 从池中移除
     */
    public void removeFromInstancePool(ConsumerInstance consumerInstance) {
        synchronized (this) {
            String topic = consumerInstance.getTopic();
            List<ConsumerInstance> currentConsumeInstanceList = consumeInstanceMap.get(topic);
            List<ConsumerInstance> filterConsumeInstanceList = currentConsumeInstanceList.stream()
                    .filter(item -> !item.getConsumerReqId().equals(consumerInstance.getConsumerReqId())).collect(Collectors.toList());
            consumeInstanceMap.put(topic, filterConsumeInstanceList);
            Set<String> consumerGroupSet = reBalanceInfo.getChangeConsumerGroupMap().get(topic);
            if (CollectionUtils.isEmpty(consumerGroupSet)) {
                return;
            }
            consumerGroupSet.remove(consumerInstance.getConsumeGroup());
        }
    }

    /**
     * 执行重平衡逻辑
     * 定时任务触发，把已有的队列分配给消费者
     */
    public void doReBalance() {
        synchronized (this) {
            String reBalanceStrategy = CommonCache.getGlobalProperties().getReBalanceStrategy();
            // 触发重平衡行为，根据参数决定重平衡策略的不同
            reBalanceInfo.setConsumeInstanceMap(this.consumeInstanceMap);
            reBalanceStrategyMap.get(reBalanceStrategy).doReBalance(reBalanceInfo);
            reBalanceInfo.getChangeConsumerGroupMap().clear();;
        }
    }

    /**
     * 启动重平衡定时任务
     */
    public void startReBalanceJob() {
        Thread reBalanceTask = new Thread(()->{
            log.info("启动重平衡任务");
            while (true) {
                try {
                    TimeUnit.SECONDS.sleep(10);
                    doReBalance();
                } catch (InterruptedException e) {
                    log.error("重平衡任务出错", e);
                }
            }
        });
        reBalanceTask.setName("重平衡任务");
        reBalanceTask.start();
    }

}
